import os
import glob
import argparse
import gmalg

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    args = parser.parse_args()
    directory = args.directory

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(10):
        Ida_file = os.path.join(directory, f'IDa_{i}')
        dA_file = os.path.join(directory, f'dA_{i}')
        PA_file = os.path.join(directory, f'PA_{i}')
        rA_file = os.path.join(directory, f'Pri_A_{i}')
        RA_file = os.path.join(directory, f'Pub_A_{i}')
        KA_file = os.path.join(directory, f'KA_{i}')
        PB_file = os.path.join(directory, f'PB_{i}')
        RB_file = os.path.join(directory, f'Pub_B_{i}')
        Idb_file = os.path.join(directory, f'IDb_{i}')
        
        # 读取文件内容
        Ida_content = read_binary_file(Ida_file)
        dA_content = read_binary_file(dA_file)
        PA_content = read_binary_file(PA_file)
        rA_content = read_binary_file(rA_file)
        RA_content = read_binary_file(RA_file)
        KA_content = read_binary_file(KA_file)
        PB_content = read_binary_file(PB_file)
        RB_content = read_binary_file(RB_file)
        Idb_content = read_binary_file(Idb_file)

        length = int(len(KA_content) / 2)
        
        # 添加到合并内容
        merged_content += f"公钥= {PA_content}\n"
        merged_content += f"私钥= {dA_content}\n"
        merged_content += f"临时公钥= {RA_content}\n"
        merged_content += f"临时私钥= {rA_content}\n"
        merged_content += f"对方公钥= {PB_content}\n"
        merged_content += f"对方临时公钥= {RB_content}\n"
        merged_content += f"ID= {Ida_content}\n"
        merged_content += f"对方ID= {Idb_content}\n"
        merged_content += f"会话密钥长度= {length:08x}\n"
        merged_content += f"会话密钥= {KA_content}\n"
        merged_content += "\n"
    
    # 写入到txt文件
    output_file = 'sm2密钥协商(发起者).txt'
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")
    
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(10):
        Ida_file = os.path.join(directory, f'IDb_{i}')
        dA_file = os.path.join(directory, f'dB_{i}')
        PA_file = os.path.join(directory, f'PB_{i}')
        rA_file = os.path.join(directory, f'Pri_B_{i}')
        RA_file = os.path.join(directory, f'Pub_B_{i}')
        KA_file = os.path.join(directory, f'KB_{i}')
        PB_file = os.path.join(directory, f'PA_{i}')
        RB_file = os.path.join(directory, f'Pub_A_{i}')
        Idb_file = os.path.join(directory, f'IDa_{i}')
        
        # 读取文件内容
        Ida_content = read_binary_file(Ida_file)
        dA_content = read_binary_file(dA_file)
        PA_content = read_binary_file(PA_file)
        rA_content = read_binary_file(rA_file)
        RA_content = read_binary_file(RA_file)
        KA_content = read_binary_file(KA_file)
        PB_content = read_binary_file(PB_file)
        RB_content = read_binary_file(RB_file)
        Idb_content = read_binary_file(Idb_file)

        length = int(len(KA_content) / 2)
        
        # 添加到合并内容
        merged_content += f"公钥= {PA_content}\n"
        merged_content += f"私钥= {dA_content}\n"
        merged_content += f"临时公钥= {RA_content}\n"
        merged_content += f"临时私钥= {rA_content}\n"
        merged_content += f"对方公钥= {PB_content}\n"
        merged_content += f"对方临时公钥= {RB_content}\n"
        merged_content += f"ID= {Idb_content}\n"
        merged_content += f"对方ID= {Ida_content}\n"
        merged_content += f"会话密钥长度= {length:08x}\n"
        merged_content += f"会话密钥= {KA_content}\n"
        merged_content += "\n"
    
    # 写入到txt文件
    output_file = 'sm2密钥协商(响应者).txt'
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")


if __name__ == "__main__":
    main()
